import torch
import cv2 as cv
from C3D_model import C3D
from utils import image_plot_conf

print('c3d.')
use_cuda = True
device = torch.device('cuda:0' if use_cuda else 'cpu')
model = C3D()
model.load_state_dict(torch.load('c3d.pth'))
model.eval()
model.to(device)
cap = cv.VideoCapture('E:/dataset/ucf101/JumpingJack/v_JumpingJack_g19_c02.avi')
frame_len = 31
clip = torch.zeros([frame_len, 240, 320, 3])
print('Network loading complete.')
cnt = 0
while cap.isOpened():
    ret, image_origin = cap.read()
    if not ret:
        break
    image = torch.tensor(image_origin/256, dtype=torch.float32)
    clip = torch.cat([clip[1:frame_len], image.unsqueeze(0)], dim=0)
    inputs = clip.unsqueeze(0).permute(0, 4, 1, 2, 3).to(device)
    predict = model(inputs)
    cnt += 1
    if cnt >= frame_len:
        image = image_plot_conf(image_origin, predict)
        cv.imshow('result', image)
        cv.waitKey(10)
